#include <cassert>
#include <stdexcept>
#include <chrono>
#include "cuda_runtime.h"
#include "cuda_runtime_api.h"
#include "cublas_v2.h"
#include "cusolverDn.h"
#include "device_launch_parameters.h"
// #include "common.h"
#include <omp.h>
// #include <atomic>
#include <cmath>
// #include "mkl.h"
#ifdef _WIN64
#include <windows.h>
#endif

// CUDA API error checking
#define CUDA_CHECK(err)                                                                            \
    do {                                                                                           \
        cudaError_t err_ = (err);                                                                  \
        if (err_ != cudaSuccess) {                                                                 \
            printf("CUDA error %d at %s:%d\n", err_, __FILE__, __LINE__);                          \
            throw std::runtime_error("CUDA error");                                                \
        }                                                                                          \
    } while (0)

// cusolver API error checking
#define CUSOLVER_CHECK(err)                                                                        \
    do {                                                                                           \
        cusolverStatus_t err_ = (err);                                                             \
        if (err_ != CUSOLVER_STATUS_SUCCESS) {                                                     \
            printf("cusolver error %d at %s:%d\n", err_, __FILE__, __LINE__);                      \
            throw std::runtime_error("cusolver error");                                            \
        }                                                                                          \
    } while (0)

// cublas API error checking
#define CUBLAS_CHECK(err)                                                                          \
    do {                                                                                           \
        cublasStatus_t err_ = (err);                                                               \
        if (err_ != CUBLAS_STATUS_SUCCESS) {                                                       \
            printf("cublas error %d at %s:%d\n", err_, __FILE__, __LINE__);                        \
            throw std::runtime_error("cublas error");                                              \
        }                                                                                          \
    } while (0)

void setProcessRealTime()
{
    int max_nt = omp_get_max_threads();
    printf("Max CPU threads: %d\n", max_nt);
    // gnt = max_nt;

#ifdef _WIN64
    bool ret;
    DWORD_PTR proc_mask = 0, sys_mask;
    for (int i = 0, bit = 0b1; i < max_nt; i++, bit = bit << 1)
        proc_mask |= bit;
    ret = SetProcessAffinityMask(GetCurrentProcess(), proc_mask);
    if (ret == 0) {
        printf("Error! SetProcessAffinityMask return %d\n", ret);
        exit(-1);
    }
    ret = GetProcessAffinityMask(GetCurrentProcess(), &proc_mask, &sys_mask);
    if (ret == 0) {
        printf("Error! GetProcessAffinityMask return %d\n", ret);
        exit(-1);
    }
    printf("ProcessAffinityMask proc %llu system %llu\n", proc_mask, sys_mask);

    DWORD prio;
    prio = GetPriorityClass(GetCurrentProcess());
    printf("Process Priority Class 0x%08x\n", prio);
    ret = SetPriorityClass(GetCurrentProcess(), 0x00000100);
    if (ret == 0) {
        printf("Error! SetPriorityClass return %d\n", ret);
        exit(-1);
    }
    prio = GetPriorityClass(GetCurrentProcess());
    printf("Process Priority Class 0x%08x\n", prio);
#endif
}

int main(int argc, char* argv[])
{
    const int N = atoi(argv[1]);
    const int test_cnt = atoi(argv[2]);
    const int real_N = atoi(argv[3]);
    printf("N %d test %d times real_N %d\n", N, test_cnt, real_N);
    // setProcessRealTime();

    double* h_A = new double [N * N], * h_x = new double [N], * h_y = new double [N], * h_B = new double [N * N];
    // printf("bp 0\n");

    double* d_A = nullptr, * d_x = nullptr, * d_y = nullptr, * d_B = nullptr;
    CUDA_CHECK(cudaMalloc((void**)&d_A, sizeof(*d_A) * N * N));
    CUDA_CHECK(cudaMalloc((void**)&d_B, sizeof(*d_B) * N * N));
    CUDA_CHECK(cudaMalloc((void**)&d_x, sizeof(*d_x) * N));
    CUDA_CHECK(cudaMalloc((void**)&d_y, sizeof(*d_y) * N));
    // printf("bp 1\n");

    srand(1998);
    for (size_t i = 0; i < N * N; i++) {
        h_A[i] = (double) rand() / RAND_MAX;
        h_B[i] = (double) rand() / RAND_MAX;
    }
    // printf("bp 1.25\n");
    for (size_t i = 0; i < N; i++) {
        h_x[i] = (double) rand() / RAND_MAX;
    }
    // printf("bp 1.5\n");
    CUDA_CHECK(cudaMemcpy(d_A, h_A, sizeof(*h_A) * N * N, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_B, h_B, sizeof(*h_B) * N * N, cudaMemcpyHostToDevice));
    CUDA_CHECK(cudaMemcpy(d_x, h_x, sizeof(*h_x) * N    , cudaMemcpyHostToDevice));
    // printf("bp 2\n");

    cublasHandle_t handle;
    cublasCreate(&handle);
    cudaEvent_t start, stop;
    CUDA_CHECK(cudaEventCreate(&start));
    CUDA_CHECK(cudaEventCreate(&stop ));
    cublasPointerMode_t mode;
    CUBLAS_CHECK(cublasGetPointerMode(handle, &mode));
    printf("mode %d\n", mode);
    CUBLAS_CHECK(cublasSetPointerMode(handle, CUBLAS_POINTER_MODE_HOST));
    double alpha = 1.0, beta = 0.0;
    float tms;

    for (int i = 0; i < 5; i++) {// warm up
        CUBLAS_CHECK(cublasDgemv(handle, CUBLAS_OP_N, real_N, real_N, &alpha, d_A, N, d_x, 1, &beta, d_y, 1));
    }
    
    CUDA_CHECK(cudaEventRecord(start));
    for (int i = 0; i < test_cnt; i++) {// test 
        CUBLAS_CHECK(cublasDgemv(handle, CUBLAS_OP_N, real_N, real_N, &alpha, d_A, N, d_x, 1, &beta, d_y, 1));
    }
    CUDA_CHECK(cudaEventRecord(stop));
    CUDA_CHECK(cudaEventSynchronize(stop));
    CUDA_CHECK(cudaEventElapsedTime(&tms, start, stop));
    tms /= test_cnt;

    CUDA_CHECK(cudaMemcpy(h_y, d_y, sizeof(*h_y) * N    , cudaMemcpyDeviceToHost));
    double ck_sum = 0.0;
    for (size_t i = 0; i < real_N; i++) {
        ck_sum += h_y[i] * (i % 10);
    }
    size_t bytes = real_N * (real_N + 2) * sizeof(*d_A);
    double BW = (double) bytes / (1024*1024*1024* (tms / 1000));
    printf("check sum %.20e Time %.6f s BW %.2f GB/s\n", ck_sum, tms / 1000, BW);

    // for (int i = 0; i < 5; i++) {
    //     cblas_dgemv(CblasColMajor, CblasNoTrans, real_N, real_N, alpha, h_A, N, h_x, 1, beta, h_y, 1);
    // }
    // auto beg = std::chrono::high_resolution_clock::now();
    // for (int i = 0; i < test_cnt; i++) {
    //     cblas_dgemv(CblasColMajor, CblasNoTrans, real_N, real_N, alpha, h_A, N, h_x, 1, beta, h_y, 1);
    // }
    // ck_sum = 0.0;
    // for (size_t i = 0; i < real_N; i++) {
    //     ck_sum += h_y[i] * (i % 10);
    // }
    // auto end = std::chrono::high_resolution_clock::now();
    // auto elapsed = (std::chrono::duration_cast<std::chrono::nanoseconds>(end - beg)).count() * 1e-9;
    // elapsed /= test_cnt;
    // BW = (double) bytes / (1024*1024*1024* elapsed);
    // printf("check sum %.20e Time %.6f s BW %.2f GB/s\n", ck_sum, elapsed, BW);

    // test trsv
    // forward: A r = y
    // backward: B x = r
    for (int i = 0; i < 5; i++) {// warm up
        CUBLAS_CHECK(cublasDtrsv(handle, CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_N, CUBLAS_DIAG_UNIT    , real_N, d_A, N, d_y, 1));
        CUBLAS_CHECK(cublasDtrsv(handle, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT, real_N, d_A, N, d_y, 1));
    }
    CUDA_CHECK(cudaEventRecord(start));
    for (int i = 0; i < test_cnt; i++) {// test 
        CUBLAS_CHECK(cublasDtrsv(handle, CUBLAS_FILL_MODE_LOWER, CUBLAS_OP_N, CUBLAS_DIAG_UNIT    , real_N, d_A, N, d_y, 1));
        CUBLAS_CHECK(cublasDtrsv(handle, CUBLAS_FILL_MODE_UPPER, CUBLAS_OP_N, CUBLAS_DIAG_NON_UNIT, real_N, d_A, N, d_y, 1));
    }
    CUDA_CHECK(cudaEventRecord(stop));
    CUDA_CHECK(cudaEventSynchronize(stop));
    CUDA_CHECK(cudaEventElapsedTime(&tms, start, stop));
    tms /= test_cnt;

    bytes = real_N * (real_N + 2) * sizeof(*d_A);
    BW = (double) bytes / (1024*1024*1024* (tms / 1000));
    printf("check sum %.20e Time %.6f s BW %.2f GB/s\n", ck_sum, tms / 1000, BW);


    delete[] h_A; delete[] h_x; delete[] h_y;
    cublasDestroy(handle);
    CUDA_CHECK(cudaEventDestroy(start)); CUDA_CHECK(cudaEventDestroy(stop));
    CUDA_CHECK(cudaFree(d_A)); CUDA_CHECK(cudaFree(d_x)); CUDA_CHECK(cudaFree(d_y));

    return 0;
}

